# 导入PyTorch库
import torch
import torch.nn as nn  # 神经网络模块


# 定义一个简单的全连接神经网络类
class SimpleNN(nn.Module):
    def __init__(self, input_dim=3, hidden_dim=4, output_dim=1):
        """
        初始化网络结构
        :param input_dim: 输入层维度（默认3）
        :param hidden_dim: 隐藏层维度（默认4）
        :param output_dim: 输出层维度（默认1）
        """
        super(SimpleNN, self).__init__()  # 调用父类nn.Module的初始化方法

        # 第一全连接层（输入层→隐藏层）
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        # 激活函数ReLU（解决梯度消失问题）
        self.relu = nn.ReLU()
        # 第二全连接层（隐藏层→输出层）
        self.fc2 = nn.Linear(hidden_dim, output_dim)
        # 输出层激活函数Sigmoid（将输出压缩到0-1，适用于二分类）
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        """
        前向传播过程
        :param x: 输入张量，形状应为[batch_size, input_dim]
        :return: 预测输出
        """
        # 第一层线性变换：W1*x + b1
        z1 = self.fc1(x)
        # 第一层激活函数：ReLU(W1*x + b1)
        a1 = self.relu(z1)
        # 第二层线性变换：W2*a1 + b2
        z2 = self.fc2(a1)
        # 输出层激活：Sigmoid(W2*a1 + b2)
        y = self.sigmoid(z2)
        return y


# 实例化网络对象
model = SimpleNN(input_dim=3, hidden_dim=4, output_dim=1)

# 创建模拟输入数据（2个样本，每个样本3个特征）
# 形状说明：[batch_size=2, input_dim=3]
x = torch.tensor([[1.0, 2.0, 3.0], [0.5, -1.0, 0.0]])

# 执行前向传播计算
output = model(x)

# 打印预测输出（每个样本的输出概率）
print("预测输出:\n", output)